import os
import pandas as pd
from .base import BaseDataset
import pickle

class MyRopeGlassBigModel(BaseDataset):
    num_classes = 5
    classnames = [
        '__background__', 'rope', 'glass','pet', 'human'
    ]
    palette = [
        (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0),(128,128,128)
    ]
    assert num_classes == len(classnames) and num_classes == len(palette)
    def __init__(self, mode, logger_handle, dataset_cfg):
        super(MyRopeGlassBigModel, self).__init__(mode=mode, logger_handle=logger_handle, dataset_cfg=dataset_cfg)
        # obtain the dirs
        rootdir = dataset_cfg['rootdir']
        self.image_dir = os.path.join(rootdir, 'JPEGImages')
        self.ann_dir = os.path.join(rootdir, 'SegmentationClass')
        self.bigmodel_dir = os.path.join(rootdir,'BigModelOutput')
        self.bigmodel_ext = '.pkl'
        self.set_dir = os.path.join(rootdir, 'ImageSets', 'Segmentation')
        # obatin imageids
        df = pd.read_csv(os.path.join(self.set_dir, dataset_cfg['set']+'.txt'), names=['imageids'])
        self.imageids = df['imageids'].values
        self.imageids = [str(_id) for _id in self.imageids]
    def __getitem__(self, index):
        # imageid
        imageid = self.imageids[index % len(self.imageids)]
        # read sample_meta
        imagepath = os.path.join(self.image_dir, f'{imageid}{self.image_ext}')
        annpath = os.path.join(self.ann_dir, f'{imageid}{self.ann_ext}')
        bigmodelpath = os.path.join(self.bigmodel_dir, f'{imageid}{self.bigmodel_ext}')
        with open(bigmodelpath, 'rb') as file:
            bigmodeloutput = pickle.load(file).to('cpu')
        sample_meta = self.read(imagepath, annpath)
        # add image id
        sample_meta.update({'id': imageid})
        # synctransforms
        sample_meta = self.synctransforms(sample_meta)
        # add bigmodel output
        sample_meta.update({'bigmodeloutput': bigmodeloutput})
        # return
        return sample_meta
